package nsfs

import (
	"GLibHac/crypto"
	"GLibHac/fssystem"
	"GLibHac/tools"
	"fmt"
	"github.com/pkg/errors"
)

const DigestSize = 0x20

var HashErr = fmt.Errorf("hash error")

type IntegrityVerificationStorage struct {
	*SectorStorage
	hashStorage         fssystem.IStorage
	integrityCheckLevel IntegrityCheckLevel
	IntegrityStorageType
	BlockValidities []Validity
}

func (i *IntegrityVerificationStorage) Read(offset int64, destination []byte) (int, error) {
	destLen := int32(len(destination))
	if destLen < 0 || destLen > i.SectorSize {
		return 0, errors.Errorf("length %d is invalid, must be between 0 and %d", destLen, i.SectorSize)
	}

	blockIndex := int32(offset / int64(i.SectorSize))

	if i.BlockValidities[blockIndex] == ValidityInvalid && i.integrityCheckLevel == IntegrityCheckLevelErrorOnInvalid {
		// TODO: Differentiate between the top and lower layers
		return 0, HashErr
	}

	needsHashCheck := i.integrityCheckLevel != IntegrityCheckLevelNone && i.BlockValidities[blockIndex] == ValidityUnchecked

	if i.IntegrityStorageType != IntegrityStorageTypeSave && !needsHashCheck {
		return i.SectorStorage.Read(offset, destination)
	}

	hashBuffer := make([]byte, DigestSize)
	hashPos := int64(blockIndex * DigestSize)
	if n, err := i.hashStorage.Read(hashPos, hashBuffer); err != nil {
		return n, err
	}

	if i.IntegrityStorageType == IntegrityStorageTypeSave {
		if tools.IsZeroArray(hashBuffer) {
			i.BlockValidities[blockIndex] = ValidityValid
			return 0, nil
		}
		if !needsHashCheck {
			return i.SectorStorage.Read(offset, destination)
		}
	}

	dataBuffer := make([]byte, i.SectorSize)
	if n, err := i.SectorStorage.Read(offset, destination); err != nil {
		return n, err
	}
	if i.BlockValidities[blockIndex] != ValidityUnchecked {
		return 0, nil
	}

	bytesToHash := i.SectorSize

	if destLen < i.SectorSize {
		// Pad out unused portion of block
		for t := destLen; t < i.SectorSize-destLen; t++ {
			dataBuffer[t] = 0
		}
		// Partition FS hashes don't pad out an incomplete block
		if i.IntegrityStorageType == IntegrityStorageTypePartitionFs {
			bytesToHash = destLen
		}
	}

	hash := crypto.Sha256(dataBuffer[:bytesToHash])

	validity := ValidityInvalid
	if tools.ArrayEquals(hashBuffer, hash) {
		validity = ValidityValid
	}
	i.BlockValidities[blockIndex] = validity
	if validity == ValidityInvalid && i.integrityCheckLevel == IntegrityCheckLevelErrorOnInvalid {
		return 0, HashErr
	}
	return 0, nil
}

func NewIntegrityVerificationStorage(info *IntegrityVerificationInfo, hashStorage fssystem.IStorage, integrityCheckLevel IntegrityCheckLevel) *IntegrityVerificationStorage {
	storage := &IntegrityVerificationStorage{
		hashStorage:          hashStorage,
		integrityCheckLevel:  integrityCheckLevel,
		IntegrityStorageType: info.IntegrityStorageType,
		SectorStorage:        NewSectorStorage(info.Data, info.BlockSize),
	}
	storage.ReadFunc = storage.Read
	storage.SizeFunc = storage.SectorStorage.Size
	storage.BlockValidities = make([]Validity, storage.SectorSize)
	return storage
}
